from libatomism import *
from math import *
import array 
import matplotlib.pyplot as plt

msLogger.setPriorityMin(Priority.ERROR)

unit=msUnitsManager.New("Angstrom Degree amu kcal/mol ps")

""" Let's define the H2O ZMatrix:
O
H 0 0.96
H 0 0.96 1 104.
"""
zmat=msZMat.New(unit).setId("zmat").set("O H 0 0.96 H 0 0.96 1 104.")
system=msSystem.New(unit).addEntity(zmat)

""" If you have VTK and if you linked atomism with it, you can display the system
by uncommenting the two next lines
"""
#viewer = msVisualizerAtoms.New()
#viewer.watch(system)

""" Rotational motion, H2O is C2v, hence the symmetry number of 2
"""
rotor=msRigidRotor.New(unit).setMmtOfInertia(system).setSymmetryNumber(2).setId("symRotor")

""" Translational motion """
trans=msTranslation.New(unit)
trans.setSystem(system)

""" Create the HO assembly 
"""
oscillsHO = msOscillatorsAssembly.New(unit)
oscillsHO.getParameters().setParameter("Emax",100/3.,unit)
oscillsHO.addOscillator( msHarmonicOscillator.New("cm-1").set(1689.00) )
oscillsHO.addOscillator( msHarmonicOscillator.New("cm-1").set(3828.44) )
oscillsHO.addOscillator( msHarmonicOscillator.New("cm-1").set(3934.84) )
oscillsHO.computeDOS()
dos2=oscillsHO.DOS()
x2=dos2.getXSample1d(0,10000)[:]
sos2=oscillsHO.SOS().getYSample1d(0,10000)[:]

""" Create the AHO assembly 
"""
oscillsAHO = msOscillatorsAssembly.New(unit)
oscillsAHO.addOscillator( msAnharmonicOscillatorP2.New("cm-1").set(1689.00,"-15.88 -22.18 -29.2") )
oscillsAHO.addOscillator( msAnharmonicOscillatorP2.New("cm-1").set(3828.44,"-22.18 -46.18 -172.14") )
oscillsAHO.addOscillator( msAnharmonicOscillatorP2.New("cm-1").set(3934.84,"-29.2 -172.14 -60.09") )
oscillsAHO.getParameters().setParameter("Emax",100/3,unit)
oscillsAHO.computeDOS()
dos=oscillsAHO.DOS()
sos=oscillsAHO.SOS()

# Plot the results
fig, axes = plt.subplots(2, 2)

axes[0, 0].plot(sos.getXSample1d(0,1000),sos.getYSample1d(0,1000),x2,sos2,"green")
axes[0, 0].set_ylabel("Vibrational sum of states",fontsize=20)
axes[0, 0].set_xlabel("energy ["+unit.getEnergyStr()+"]",fontsize=20)
axes[0, 0].tick_params(labelsize=20)

Tlist=[100*i for i in range(1,30)]
QHO=[]
QAO=[]
s=[]
sHO=[]
cp=[]
cpHO=[]
for T in Tlist:
    QHO.append(oscillsHO.QfromDOS(T))   
    QAO.append(oscillsAHO.QfromDOS(T))   
    cp.append(trans.Cv(T) + rotor.Cv(T) + oscillsAHO.Cv(T) +  unit.convert("J/mol.K^-1",R) )
    cpHO.append(trans.Cv(T) + rotor.Cv(T) + oscillsHO.Cv(T) +  unit.convert("J/mol.K^-1",R) )    
    s.append(trans.S(T) + rotor.S(T) + oscillsAHO.S(T) +  unit.convert("J/mol.K^-1",R) )
    sHO.append(trans.S(T) + rotor.S(T) + oscillsHO.S(T) +  unit.convert("J/mol.K^-1",R) )    
    
axes[0, 1].plot(Tlist,QHO,"green",Tlist,QAO,"blue")
axes[0, 1].set_ylabel("Vibrational partition function",fontsize=20)
axes[0, 1].set_xlabel("T ["+unit.getTempStr()+"]",fontsize=20)
axes[0, 1].tick_params(labelsize=20)
# this is the experimental values, coming from the janaf table,  
# see http://kinetics.nist.gov/janaf/html/H-064.html, defined in J/mol/K
TListexp  =  [200,     300,     500,     800,     1200,   1500,    2000,    2500,    3000   ] 
ExpListS   = [175.5,   189.0,   206.5,   223.8,   240.5,  250.6,   264.8,   276.5,   286.5]   
ExpListCp =  [33.4,    33.6,    35.3,    38.7,    43.8,   47.1,    51.2,    53.9,    55.8 ] 


axes[1, 0].plot(Tlist,s,"blue",Tlist,sHO,"green",TListexp,[ unit.convert("J/mol",x) for x in ExpListS],"r*",markersize=15)
axes[1, 0].set_xlabel("Temperature ["+unit.getTempStr()+"]",fontsize=20)
axes[1, 0].set_ylabel("Total entropy ["+unit.getEnergyStr()+"]",fontsize=20)
axes[1, 0].tick_params(labelsize=20)

axes[1, 1].plot(Tlist,cp,"blue",Tlist,cpHO,"green",TListexp,[ unit.convert("J/mol",x) for x in ExpListCp],"r*",markersize=15)
axes[1, 1].set_xlabel("Temperature ["+unit.getTempStr()+"]",fontsize=20)
axes[1, 1].set_ylabel("Total Cp ["+unit.getEnergyStr()+"]",fontsize=20)
axes[1, 1].tick_params(labelsize=20)


fig.set_size_inches(18,18)
fig.savefig('H2O_0.png')

fig.show()

 
